import torch
import gc

# 假设你有一个大的张量或模型在GPU上
large_tensor = torch.randn(4096, 4096, device='cuda')
print(f"创建张量后，已分配显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"创建张量后，缓存显存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

# 当你不再需要这个张量时
# 1. 删除引用
del large_tensor

# 2. 执行垃圾回收
gc.collect()

# 3. 清空缓存
torch.cuda.empty_cache()

print("\n--- 清理后 ---")
print(f"清理后，已分配显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"清理后，缓存显存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
